/***************************************************************************
 *   Copyright (C) 2007 by Sassan Panahinejad (ico2)			   *
 *   ico2ico2@gmail.com							   *
 *                                                                         *
 *   This file is part of libaufs					   *
 *   Libaufs is free software; you can redistribute it and/or modify	   *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 3 of the License, or     *
 *   (at your option) any later version.				   *
 *									   *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of	   *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.			   *
 *									   *
 *   You should have received a copy of the GNU General Public License	   *
 *   along with this program.  If not, see <http://www.gnu.org/licenses/>. *
 ***************************************************************************/

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <mntent.h>


#include <sys/mount.h>

#include "aufs.h"
#include "error.h"


//move to aufs.h
#define AUFS_FSTYPE "aufs"
#define MTAB "/proc/mounts"

#define BUFFER_SIZE 4096

/*
        Adds a branch to an aufs union.
	If unsure, use either of the 2 wrapper functions provided.
        Argument 1 is the path of the union to modify.
        Argument 2 is the path to add to the union.
        Argument 3 specifies whether or not the branch should be writable, 1=rw, 0=ro
        Argument 4 is the index to add the branch.
        Return value is 0 on success, <0 on failure.
*/
int aufs_add_branch(char *unionpath, char *branchpath, int writable, int index)
{
        int ret;
        char *mount_options, *writestring;
	if (writable)
		writestring = "rw";
	else
		writestring = "ro";
        if (index >= 0) 
                asprintf(&mount_options, "ins:%d:%s=%s", index, branchpath, writestring);
        else
                asprintf(&mount_options, "append:%s=%s", branchpath, writestring);
        debug_output("Attempting to add to aufs union using mount options: %s\n", mount_options);
        ret = mount("aufs", unionpath, "aufs", MS_REMOUNT, mount_options);
        if (ret < 0) {
                error_output("Failed to add branch %s to aufs union %s: %m\n", branchpath, unionpath);
                return -1;
        }
        debug_output("Successfully added branch %s to aufs union %s\n", branchpath, unionpath);
        return 0;
}



/*
	Adds a branch before a specified branch.
*/
int aufs_add_branch_before(char *unionpath, char *branchpath, int writable, char *before)
{
	int index;
	index = aufs_get_index(unionpath, before);
	if (index < 0) {
		info_output("Unable to get index\n");
		return -1;
	}
	return aufs_add_branch(unionpath, branchpath, writable, index);
}

/*
	Adds a branch after a specified branch.
*/
int aufs_add_branch_after(char *unionpath, char *branchpath, int writable, char *after)
{
	int index;
	index = aufs_get_index(unionpath, after)+1;
	if (index < 0) {
		info_output("Unable to get index\n");
		return -1;
	}
	return aufs_add_branch(unionpath, branchpath, writable, index);
}

/*
        Removes a branch from an aufs union.
        Argument 1 is the path of the union to modify.
        Argument 2 is the path of the branch to remove.
        Returns 0 on success, <0 on failure.
*/
int aufs_remove_branch(char *unionpath, char *branchpath)
{
        int ret;
        char *mount_options;
        asprintf(&mount_options, "remount,del:%s", branchpath);
        debug_output("Attempting to remove from aufs union using mount options: %s\n", mount_options);
        ret = mount("aufs", unionpath, "aufs", 0, mount_options);
        if (ret < 0) {
                error_output("Failed to remove branch %s from aufs union %s: %m\n", branchpath, unionpath);
                return -1;
        }
        debug_output("Successfully removed branch %s from aufs union %s: %m\n", branchpath, unionpath);
        return 0;
}

/*
        Changes the permissions of an aufs union.
        Argument 1 is the path of the union to modify.
        Argument 2 is the path of the branch to update.
        Argument 3 specifies whether or not the path should be writable, 1=rw, 0=ro
        Return value is 0 on success, <0 on failure.
*/
int aufs_set_perms(char *unionpath, char *branchpath, int writable)
{
        int ret;
        char *mount_options;
	char *perms;
	if (writable)
		perms = "rw";
	else
		perms = "ro";
        asprintf(&mount_options, "remount,mod:%s=%s", branchpath, perms);
        debug_output("Attempting to set permisions on aufs branch %s in union %s to %s\n", branchpath, unionpath, perms);
        ret = mount("aufs", unionpath, "aufs", 0, mount_options);
        if (ret < 0) {
                error_output("Failed to set permissions on branch %s from aufs union %s to %s: %m\n", branchpath, unionpath, perms);
                return -1;
        }
        debug_output("Set permisions on aufs branch %s in union %s to %s\n", branchpath, unionpath, perms);
        return 0;
	
}

/*
        Checks whether or not a union can be written to.
        Argument 1 is the path of the union to read.
        Argument 2 is the path of the branch to check.
        Return value 1=rw, 0=ro, <0=error
*/
int aufs_is_writable(char *unionpath, char *branchpath)
{
	char *ptr, *branchstring;
	debug_output("Trying to determine if branch %s of union %s is writable\n", branchpath, unionpath);
	branchstring = aufs_get_branchstring(unionpath);
	ptr = strstr(branchstring, branchpath);
	if (!ptr)
		return -1;
	ptr += strlen(branchpath);
	if ( ptr[0] != '=' || ptr[2] == 'w') {
		debug_output("Branch %s of union %s is writable\n", branchpath, unionpath);
		return 1;
	}
	debug_output("Branch %s of union %s is not writable\n", branchpath, unionpath);
	return 0;	
	
}

/*
        Retrieves the the branches string from mtab (possibly brs also?)
        Argument 1 is the path of the union to read.
        Return value is the string. NULL on failure.
*/
char *aufs_get_branchstring(char *unionpath)
{
	struct mntent *current;
	FILE *mtab;
	char *branchstring, *branchstr, *realunionpath;
	debug_output("Attempting to retrieve branchstring for union %s\n", unionpath);
	realunionpath = realpath(unionpath, NULL);
	mtab = setmntent(MTAB, "r");
	while ((current = getmntent(mtab))) {
		debug_output("looking for union in dir: %s\n", current->mnt_dir);
		if (!strcmp(current->mnt_dir, realunionpath)) {
			branchstring =  strstr(current->mnt_opts, "br:");
			if (!branchstring) {
				info_output("Cannot extract branch info from mount options\n");
				debug_output("mount options for %s (%s): %s\n", current->mnt_dir, current->mnt_fsname,
						 current->mnt_opts);
				continue;
			}
			branchstring += strlen("br:");
			endmntent(mtab);
			branchstr = strdup(branchstring);
			debug_output("Branchstring for union %s is %s\n", unionpath, branchstr);
			return branchstr;
		}
	}
	endmntent(mtab);
	debug_output("Unable to retrieve branchstring for union %s\n", unionpath);
	return NULL;
}

/*
        Retrieves the first branch of a union.
        Argument 1 is the path of the union to modify.
        Argument 2 is the memory location to store the branch information string.
        Memory is allocated and the string is stored in here.
        Return value is the path of the branch.
*/
char *aufs_first_branch(char *unionpath, char **branchstring)
{
	*branchstring = aufs_get_branchstring(unionpath);
	return aufs_next_branch(branchstring);
	
}

/*
        Retrieves the next branch of a union.
        Argument 1 is the branch information string.
        This value is modified and memory freed appropriatly.
        Return value is the path of the branch. NULL if not found
*/
char *aufs_next_branch(char **branchstring)
{
	char *tmp, *new, *ptr;
	debug_output("Attempting to retrieve a branch\n");
	tmp = new = *branchstring;
	if (!branchstring || !*branchstring) {
		return NULL;
	}
	new = strchr(new, ':');
	if (!new)
		*branchstring = NULL;
	else {
		new++[0] = '\0';
		*branchstring = new;
	}
	ptr = strchr(tmp, '=');
	if (!ptr)
		ptr = tmp + strlen(tmp)-1;
	ptr[0] = '\0';
	return tmp;
		
}

/*
        Returns the index of a branch in an aufs union.
        Argument 1 is the path of the union to read.
        Argument 2 is the path of the branch to check.
        Return value is the index, -1 if not found.
*/
int aufs_get_index(char *unionpath, char *branchpath)
{
        char *branch, *branchstring, *realbranchpath;
        int index = 0;
	realbranchpath = realpath(branchpath, NULL);
	debug_output("Attempting to get index of branch %s in union %s\n", realbranchpath, unionpath);
        for (branch = aufs_first_branch(unionpath, &branchstring); branch;
                        branch = aufs_next_branch(&branchstring)) {
                if (!strcmp(branch, realbranchpath)) {
			debug_output("Index of branch %s in union %s is %d\n", 
					realbranchpath, unionpath, index);
                        return index;
		}
		index++;
        }
        return -1;

}

/*
	Creates an aufs union.
	Argument is the point to mount the union.
	Return value is 0 on sucess, <0 on failure.
*/
int aufs_create_union(char *unionpath, char *branchpath, int writable)
{
	int ret;
	char *mount_options, *writestring;
	if (writable)
		writestring = "rw";
	else
		writestring = "ro";
	asprintf(&mount_options, "dirs=%s=%s", branchpath, writestring);
	ret = mount("aufs", unionpath, "aufs", 0, mount_options);
	if (ret < 0) {
		error_output("Unable to create union at %s: %m\n", unionpath);
		return -1;
	}
	info_output("Created union at %s\n", unionpath);
	return 0;
}

/*
	Checks if a directory contains an aufs union.
	Cheats and uses aufs_get_branchstring which is wasteful.
	Argument is the directory to check for union.
	return value is 1 if the directory contains a union, 0 if it does not.
*/
int aufs_is_union(char *dir)
{
	if (aufs_get_branchstring(dir))
		return 1;
	else
		return 0;
}

